#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
文本转语音（TTS）ROS节点
该节点订阅 'tts_text' 话题上的字符串消息，使用火山引擎的语音合成服务
将文本转换为PCM音频流，并通过pyaudio实时播放。

核心功能:
- 订阅ROS话题获取待合成文本。
- 通过WebSocket与火山TTS服务进行二进制通信。
- 使用独立的线程管理音频播放，避免阻塞ROS回调。
- 实现了一个生产者-消费者模型：WebSocket接收音频数据（生产者），
  音频播放线程消费数据（消费者）。
- 支持播放中断：当接收到新的TTS请求时，会清空当前待播放的音频队列，
  从而实现即时打断当前播报。
"""

import rospy
import asyncio
import websockets
import uuid
import json
import threading
import struct
import gzip
import copy
from std_msgs.msg import String

import pyaudio
import queue
import time
import base64


class TTSNode:
    """
    封装了TTS功能的ROS节点类。
    管理与TTS服务的WebSocket连接、音频数据流处理和本地音频播放。
    """
    def __init__(self):
        """
        初始化TTS节点。
        """
        rospy.init_node('tts_node', anonymous=True)
        
        # 订阅 'tts_text' 话题，接收需要转换为语音的文本
        rospy.Subscriber('tts_text', String, self.tts_callback)
        
        # 发布TTS状态（"playing", "finished"），供其他节点了解播报是否正在进行
        self.tts_status_pub = rospy.Publisher('tts_status', String, queue_size=10)
        
        # 从ROS参数服务器或环境变量中获取鉴权信息和配置
        self.app_id = rospy.get_param('~app_id')
        self.access_key = rospy.get_param('~access_key')
        self.voice_type = rospy.get_param('~voice_type', 'zh_male_shaonianzixin_moon_bigtts')
        
        # 火山引擎语音合成服务的WebSocket接口地址
        self.api_url = 'wss://openspeech.bytedance.com/api/v1/tts/ws_binary'
        
        # 音频播放参数，使用16-bit PCM格式
        self.format = pyaudio.paInt16
        self.channels = 1              # 单声道
        self.rate = 16000              # 默认采样率
        self.chunk = 2048              # 每次读写的音频帧数
        
        # 与TTS服务通信的二进制协议头
        self.default_header = bytearray(b'\x11\x10\x11\x00')
        
        # 消息类型定义，用于解析服务器响应
        self.MESSAGE_TYPES = {11: "audio-only server response", 12: "frontend server response", 15: "error message from server"}
        self.MESSAGE_TYPE_SPECIFIC_FLAGS = {0: "no sequence number", 1: "sequence number > 0",
                                           2: "last message from server (seq < 0)", 3: "sequence number < 0"}
        
        # 初始化pyaudio，用于音频播放
        self.audio = pyaudio.PyAudio()
        self.stream = None
        self.audio_lock = threading.Lock()  # 锁，用于保护对音频流的并发访问
        
        # 创建一个线程安全的队列，用作音频数据的缓冲区
        self.audio_queue = queue.Queue(maxsize=100)
        
        # 播放状态管理
        self.is_playing = False  # 标志是否正在播放音频
        self.tts_finished = True # 标志WebSocket数据流是否接收完毕
        
        # 创建并启动音频播放后台线程
        self.play_thread = threading.Thread(target=self.audio_playback_thread)
        self.play_thread.daemon = True
        self.play_thread.start()
        
        rospy.loginfo("语音合成节点已启动")
        rospy.loginfo(f"当前使用的音色: {self.voice_type}")

    def init_audio_stream(self):
        """
        初始化或重新初始化pyaudio音频流。
        会自动搜索可用的输出设备，并尝试使用兼容的采样率。
        """
        with self.audio_lock:
            try:
                # 如果存在旧的流，先安全关闭
                if self.stream and not self.stream.is_stopped():
                    self.stream.stop_stream()
                if self.stream:
                    self.stream.close()
                
                # 智能选择音频输出设备，优先选择'pulse'或'default'
                output_device = None
                for i in range(self.audio.get_device_count()):
                    device_info = self.audio.get_device_info_by_index(i)
                    if device_info['maxOutputChannels'] > 0:
                        if 'pulse' in device_info['name'].lower():
                            output_device = i
                            break
                        elif 'default' in device_info['name'].lower() and output_device is None:
                            output_device = i
                
                if output_device is None:
                    output_device = self.audio.get_default_output_device_info()['index']
                
                # 尝试使用不同的采样率打开音频流，以提高兼容性
                supported_rates = [16000, 22050, 44100, 48000]
                for rate in supported_rates:
                    try:
                        self.stream = self.audio.open(
                            format=self.format,
                            channels=self.channels,
                            rate=rate,
                            output=True,
                            output_device_index=output_device,
                            frames_per_buffer=self.chunk
                        )
                        # 如果成功打开，更新当前采样率并跳出循环
                        if rate != self.rate:
                            self.rate = rate
                        break
                    except Exception:
                        continue
                
                if not self.stream:
                    raise Exception("无法创建音频流，已尝试所有支持的采样率")
                    
                rospy.loginfo(f"音频流已成功初始化，设备索引: {output_device}, 采样率: {self.rate}Hz")
                return True
            except Exception as e:
                rospy.logerr(f"音频流初始化失败: {e}")
                self.stream = None
                return False

    def audio_playback_thread(self):
        """
        作为后台线程运行，持续从音频队列中取出数据并播放。
        这是一个典型的消费者线程。
        """
        self.init_audio_stream()
        
        while not rospy.is_shutdown():
            try:
                # 确保音频流是可用的
                if self.stream is None:
                    rospy.logwarn("音频流无效，尝试重新初始化...")
                    if not self.init_audio_stream():
                        time.sleep(1) # 如果初始化失败，等待一段时间再重试
                        continue

                # 从队列中获取音频数据，设置超时以避免永久阻塞
                data = self.audio_queue.get(timeout=0.1) 

                if data == b'':  # 收到中断标记
                    rospy.loginfo("收到音频中断标记，清空播放状态。")
                    if self.is_playing:
                        self.tts_status_pub.publish("finished")
                        self.is_playing = False
                    self.audio_queue.task_done() 
                    continue

                if not self.is_playing:
                    self.is_playing = True
                    self.tts_finished = False # 开始播放新一段语音，重置完成标志
                    rospy.loginfo("TTS开始播放")
                    self.tts_status_pub.publish("playing")
                
                self.stream.write(data) # 将音频数据写入流进行播放
                self.audio_queue.task_done()

            except queue.Empty:
                # 队列为空时，需要判断是否意味着一段语音已经播放完毕
                if self.is_playing and self.tts_finished:
                    rospy.loginfo("音频队列为空且TTS数据接收完毕，等待硬件缓冲区播放完成...")
                    try:
                        # 等待pyaudio的内部缓冲区完全播放完毕
                        timeout_sec = 2.0
                        start = time.time()
                        while self.stream and self.stream.is_active():
                            if rospy.is_shutdown() or (time.time() - start > timeout_sec):
                                break
                            time.sleep(0.05)
                    finally:
                        if self.is_playing:
                            rospy.loginfo("TTS播放完成，发布'finished'状态。")
                            self.tts_status_pub.publish("finished")
                            self.is_playing = False
                continue # 队列为空是正常情况，继续循环等待

            except Exception as e:
                rospy.logerr(f"音频播放线程发生严重错误: {e}")
                if self.is_playing: 
                    self.tts_status_pub.publish("finished")
                    self.is_playing = False
                
                # 发生错误后，尝试关闭并重新初始化音频流
                self.stream = None # 标记流为无效
                time.sleep(1)

    def create_request_frame(self, text):
        """
        根据输入文本构建符合火山TTS服务要求的二进制请求帧。
        
        Args:
            text (str): 需要合成的文本。
        
        Returns:
            bytes: 包含了头部、压缩后的JSON负载的完整请求数据。
        """
        # 1. 构建JSON请求体
        request_json = {
            "app": {
                "appid": self.app_id,
                "token": self.access_key,
                "cluster": "volcano_tts"
            },
            "user": {
                "uid": str(uuid.uuid4())
            },
            "audio": {
                "voice_type": self.voice_type,
                "encoding": "pcm",
                "speed_ratio": 1.0,
                "volume_ratio": 1.0,
                "pitch_ratio": 1.0,
                "rate": self.rate
            },
            "request": {
                "reqid": str(uuid.uuid4()),
                "text": text,
                "text_type": "plain",
                "operation": "submit" # 提交合成请求
            }
        }
        
        # 2. 转换为字符串并使用gzip压缩
        payload_bytes = str.encode(json.dumps(request_json))
        payload_bytes = gzip.compress(payload_bytes)
        
        # 3. 构造完整的二进制帧：固定头 + 负载长度 + 压缩负载
        full_client_request = bytearray(self.default_header)
        full_client_request.extend((len(payload_bytes)).to_bytes(4, 'big'))
        full_client_request.extend(payload_bytes)
        
        return bytes(full_client_request)

    def parse_response(self, res):
        """
        解析来自TTS服务的二进制响应。
        
        Args:
            res (bytes): WebSocket接收到的原始二进制数据。
            
        Returns:
            tuple: (消息类型, 序列号, 负载数据) 或 None
        """
        try:
            # 解析头部信息
            res_bytearray = bytearray(res)
            magic_num = res_bytearray[0]
            version = res_bytearray[1]
            header_size = res_bytearray[2]
            message_type_specific_flag = res_bytearray[3]
            serialize_type = res_bytearray[4]
            message_type = res_bytearray[5]
            
            payload_size = int.from_bytes(res_bytearray[header_size-4:header_size], 'big')
            payload = res_bytearray[header_size:header_size+payload_size]
            
            # 根据消息类型处理负载
            if message_type == 12: # 前端服务器响应
                meta_info_size = int.from_bytes(payload[0:4], 'big')
                meta_info = json.loads(gzip.decompress(payload[4:4+meta_info_size]).decode("utf-8"))
                
                # 如果是最后一条消息，设置完成标志
                if meta_info.get("data", {}).get("finish_reason", 0) == 1:
                    rospy.loginfo("收到TTS数据流结束标志。")
                    self.tts_finished = True
                
                return "meta", meta_info.get("sequence", -1), meta_info
            
            elif message_type == 11: # 纯音频响应
                sequence = int.from_bytes(payload[0:4], 'big')
                audio_data = payload[4:]
                return "audio", sequence, audio_data
            
            elif message_type == 15: # 错误消息
                rospy.logerr(f"收到服务器错误消息: {payload.decode('utf-8')}")
                self.tts_finished = True
                return "error", -1, payload.decode('utf-8')
            
            return None, None, None
            
        except Exception as e:
            rospy.logerr(f"解析服务器响应失败: {e}")
            self.tts_finished = True
            return None, None, None
    
    def process_pcm_audio(self, pcm_data):
        """
        将解码后的PCM音频数据放入播放队列。
        """
        # 简单地将音频块放入队列
        self.audio_queue.put(pcm_data)

    async def tts_request(self, text):
        """
        执行一次完整的TTS请求。
        这是一个异步函数，负责建立WebSocket连接、发送请求并接收所有响应。
        """
        self.tts_finished = False # 开始新的请求，重置完成标志
        try:
            async with websockets.connect(self.api_url) as ws:
                # 1. 发送请求帧
                request_frame = self.create_request_frame(text)
                await ws.send(request_frame)
                rospy.loginfo("TTS请求已发送，等待音频流...")
                
                # 2. 循环接收服务器响应
                while not rospy.is_shutdown():
                    res = await ws.recv()
                    msg_type, seq, payload = self.parse_response(res)
                    
                    if msg_type == "audio":
                        self.process_pcm_audio(payload)
                    elif msg_type == "meta" and payload.get("data", {}).get("finish_reason", 0) == 1:
                        # 收到结束帧，退出接收循环
                        break
                    elif msg_type == "error":
                        # 收到错误，退出
                        break
                        
        except Exception as e:
            rospy.logerr(f"TTS WebSocket通信失败: {e}")
        finally:
            # 无论成功失败，最后都将完成标志置为True
            rospy.loginfo("TTS WebSocket连接关闭。")
            self.tts_finished = True

    def tts_callback(self, msg):
        """
        ROS话题 'tts_text' 的回调函数，是所有TTS任务的入口。
        """
        text = msg.data
        if not text:
            return
        
        rospy.loginfo(f"接收到新的TTS请求: '{text}'")
        
        # --- 实现播放中断的关键 ---
        # 如果当前正在播放，清空队列以停止旧的语音
        if not self.audio_queue.empty():
            rospy.loginfo("检测到正在播放，中断当前语音。")
            with self.audio_queue.mutex:
                self.audio_queue.queue.clear()
            # 发送一个中断标记，让播放线程知道状态已改变
            self.audio_queue.put(b'')

        # 在新线程中运行asyncio事件循环，以避免阻塞回调
        def run_async_tts():
            try:
                # 获取或创建一个新的事件循环
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                # 运行tts_request协程直到完成
                loop.run_until_complete(self.tts_request(text))
                loop.close()
            except Exception as e:
                rospy.logerr(f"运行异步TTS任务时出错: {e}")

        # 启动后台线程执行异步任务
        tts_thread = threading.Thread(target=run_async_tts)
        tts_thread.start()

    def shutdown(self):
        """
        节点关闭时的清理操作。
        """
        rospy.loginfo("正在关闭TTS节点...")
        # 等待音频队列处理完毕
        self.audio_queue.join()
        
        # 安全地关闭音频流和pyaudio
        with self.audio_lock:
            if self.stream:
                self.stream.stop_stream()
                self.stream.close()
            self.stream = None
        if self.audio:
            self.audio.terminate()
        rospy.loginfo("TTS节点已关闭。")

if __name__ == '__main__':
    try:
        tts_node = TTSNode()
        rospy.on_shutdown(tts_node.shutdown) # 注册关闭钩子
        rospy.spin()
    except rospy.ROSInterruptException:
        pass